"""
日志记录工具
统一管理训练过程中的日志输出
"""
import logging
import os
from datetime import datetime


def setup_logger(log_dir, model_name, fold_idx=None):
    """
    设置日志记录器

    Args:
        log_dir: 日志保存目录
        model_name: 模型名称
        fold_idx: 折数索引（用于K折交叉验证）

    Returns:
        logger: 配置好的日志记录器
    """
    # 创建日志目录
    os.makedirs(log_dir, exist_ok=True)

    # 构建日志文件名
    if fold_idx is not None:
        log_filename = f"{model_name}_fold_{fold_idx}.log"
    else:
        log_filename = f"{model_name}.log"

    log_path = os.path.join(log_dir, log_filename)

    # 创建logger
    logger = logging.getLogger(f"{model_name}_fold_{fold_idx}")
    logger.setLevel(logging.INFO)

    # 清除已有的handlers（避免重复）
    if logger.hasHandlers():
        logger.handlers.clear()

    # 创建文件handler
    file_handler = logging.FileHandler(log_path, mode='a', encoding='utf-8')
    file_handler.setLevel(logging.INFO)

    # 创建控制台handler
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)

    # 设置日志格式
    formatter = logging.Formatter(
        '[%(asctime)s] [%(levelname)s] %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )
    file_handler.setFormatter(formatter)
    console_handler.setFormatter(formatter)

    # 添加handlers
    logger.addHandler(file_handler)
    logger.addHandler(console_handler)

    logger.info(f"=" * 80)
    logger.info(f"日志记录器初始化完成")
    logger.info(f"模型: {model_name}")
    if fold_idx is not None:
        logger.info(f"折数: {fold_idx}")
    logger.info(f"日志文件: {log_path}")
    logger.info(f"=" * 80)

    return logger


def log_config(logger, config):
    """
    记录配置信息

    Args:
        logger: 日志记录器
        config: 配置字典
    """
    logger.info("训练配置:")
    for key, value in config.items():
        logger.info(f"  {key}: {value}")
